package cn.com.bluemoon.daps.common.toolkit;

import cn.com.bluemoon.daps.common.toolkit.valueobject.MyQueryCallable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StopWatch;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.*;

/**
 * 针对数据量大的并行查询数据库工具
 *
 * @author Jarod Kong
 * @date 2021-03-27 08:52
 **/
@Slf4j
public class DbParallelHelper {

    private ExecutorService dbExecutor;

    public static DbParallelHelper build() {
        return new DbParallelHelper();
    }

    public DbParallelHelper executor(ExecutorService dbExecutor) {
        this.dbExecutor = dbExecutor;
        return this;
    }

    public ExecutorService getDbExecutor() {
        if (dbExecutor == null) {
            dbExecutor = Executors.newWorkStealingPool();
        }
        return dbExecutor;
    }

    /**
     * 提供并发查询能力
     * 提供自定义执行器的具体实现Executor -done
     * 可以考虑把查询的总数量也抽象掉（能否通过查询条件MyQueryCallable中基于查询条件拿到总数？）-done
     *
     * @param taskCallable 需要查询的逻辑+指定查询的总数，提供offset limit的入参
     * @param limit        一次查询多少
     * @return 结果集
     */
    public <T> List<T> multiDBQuery(MyQueryCallable<T> taskCallable, int limit) {
        // 边界判断
        BmAssetUtils.notNull(taskCallable, "实现对应查询的业务逻辑！！");
        int count = taskCallable.count();
        if (count <= 0 || limit <= 0) {
            return Collections.emptyList();
        }
        // count为上限
        if (count < limit) {
            limit = count;
        }
        StopWatch sw = new StopWatch();
        sw.start("组装查询并行任务[count=" + count + ",limit=" + limit + "]");
        //需要查询的次数
        int times = count / limit;
        if (count % limit != 0) {
            times = times + 1;
        }
        //开始页数
        int offset = 0;
        List<T> result = new ArrayList<>(count);
        //Callable用于产生结果
        List<Callable<List<T>>> tasks = new ArrayList<>();
        for (int i = 0; i < times; i++) {
            Callable<List<T>> tCall = taskCallable.callTask(offset, limit);
            tasks.add(tCall);
            offset = offset + limit;
            if (offset > count) {
                break;
            }
        }
        sw.stop();
        sw.start("执行器处理任务[tasks=" + tasks.size() + "]");
        //定义固定长度的线程池  防止线程过多
        ExecutorService dbExecutor = getDbExecutor();
//        ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors()+1);
        //Future用于获取结果
        List<Future<List<T>>> futures = null;
        try {
            futures = dbExecutor.invokeAll(tasks);
            //处理线程返回结果
            if (futures.size() > 0) {
                for (Future<List<T>> future : futures) {
                    result.addAll(future.get());
                }
            }
        } catch (InterruptedException | ExecutionException e) {
            e.printStackTrace();
        } finally {
            dbExecutor.shutdown();
        }
        sw.stop();
        log.info(" multi query db, size:{}, detail:{} {}", result.size(), System.lineSeparator(), sw.prettyPrint());
        return result;
    }


}
